

clear


%% data and prior
load('../../true_DGP/priorWSigma.mat')
data=readtable('../../../Data/data_main_sample.csv');

dataLT=readtable('../../../Data/data_w_treismann.csv');
dataLT.leaderturn=str2double(dataLT.leaderturn);
% fix North Vietnam
dataLT.leaderturn(dataLT.ccode==816)=0;
dataLT.leaderturn(dataLT.ccode==816&dataLT.year==1969)=1;
dataLT.leaderturn(dataLT.ccode==816&dataLT.year==1980)=1;
dataLT.leaderturn(dataLT.ccode==816&dataLT.year==1987)=1;
dataLT.leaderturn(dataLT.ccode==816&dataLT.year==1992)=1;
dataLT.leaderturn(dataLT.ccode==816&dataLT.year==1997)=1;

data=data(data.year>=1950&data.year<=2000&isfinite(data.D),:);
ccodes=unique(data.ccode);
years=unique(data.year(data.year>1950));
N=length(ccodes);
T=length(years);
K=3;
y=zeros(N,T); D=y; M=y; X=repmat(y,1,1,K);
for n=1:N
    dn=data(data.ccode==ccodes(n),:);
    dtn=dataLT(dataLT.ccode==ccodes(n),:);
    for t=1:T
        if isempty(setdiff(years(t),dn.year))
            if isfinite(dn.y(years(t)==dn.year))&&...
                    isfinite(dn.Y(years(t)-1==dn.year))&&...
                    isfinite(dtn.leaderturn(years(t)-1==dtn.year))
                    M(n,t)=1; % not missing
                    y(n,t)=dn.y(years(t)==dn.year);
                    D(n,t)=dn.D(years(t)==dn.year);
                    X(n,t,1)=log(dn.Y(years(t)-1==dn.year));
                    X(n,t,2)=dtn.leaderturn(years(t)-1==dtn.year);
                    X(n,t,3)=log(dn.Y(years(t)-1==dn.year))*dtn.leaderturn(years(t)-1==dtn.year);
            end
        end
    end
end
for k=1:K
    X(:,:,k)=M.*(X(:,:,k)-sum(sum(X(:,:,k)))/sum(sum(M==1)));
end
clear dn dtn t k

dd=readtable('../../../Data/distance_capitals.csv');
Z=10^(-6)*table2array(dd(:,2:end));
KZ=size(Z,3);

ep_sd=sqrt(diag(Sigma));

% Sparse-grid integration
[sgi_nodes,sgi_weights]=nwspgr('KPN',2,8);

cnames=cell(N,1);
for n=1:N
    a=unique(data.idacr(data.ccode==ccodes(n)));
    cnames{n}=a{1};
end

ccodes=table(ccodes,cnames);

clear a cnames data dataLT dd n


%% for Knitro

% Jacobian sparsity pattern:
JPattern=[repmat(eye(N),T,1),ones(N*T,2),ones(N*T,2*N),ones(N*T,N),...
    repmat(eye(N),T,1),zeros(N*T,N),ones(N*T,K),ones(N*T,KZ),eye(N*T)];
JPattern=sparse(JPattern);

% Hessian sparsity pattern
% log-posterior
HPattern=zeros(6*N+2+K+KZ+N*T,6*N+2+K+KZ+N*T);
  % alpha,theta,beta,v,f
HPattern(1:(5*N+2),1:(5*N+2))=eye(5*N+2);
  % vs
HPattern(5*N+2+(1:N),:)=[zeros(N,5*N+2),eye(N),zeros(N,K+KZ),repmat(eye(N),1,T)];
  % kappa
HPattern((6*N+2+K+KZ+1):end,:)=[zeros(N*T,5*N+2),repmat(eye(N),T,1),zeros(N*T,K+KZ),eye(N*T)];
% equilibrium constraints
  % alpha
  HPattern(1:N,:)=HPattern(1:N,:)+...
        [repmat(eye(N),1,2),ones(N,2*N+2),eye(N),zeros(N),ones(N,K+KZ),repmat(eye(N),1,T)];
  % theta
  HPattern(N+(1:2),:)=HPattern(N+1,:)+ones(2,6*N+2+K+KZ+N*T);
  % f
  HPattern(4*N+(1:N),:)=HPattern(4*N+(1:N),:)+...
        [repmat(eye(N),1,2),ones(N,2*N+2),eye(N),zeros(N),ones(N,K+KZ),repmat(eye(N),1,T)];
  % beta,v,xi,gamma
for n=2:4
    HPattern((n-1)*N+2+(1:N),:)=HPattern((n-1)*N+2+(1:N),:)+...
        [ones(N,5*N+2),zeros(N),ones(N,K+KZ+N*T)];
end
HPattern(6*N+2+(1:K),:)=HPattern(6*N+2+(1:K),:)+...
        [ones(K,5*N+2),zeros(K,N),ones(K,K+KZ+N*T)];
HPattern(6*N+2+K+(1:KZ),:)=HPattern(6*N+2+K+(1:KZ),:)+...
        [ones(KZ,5*N+2),zeros(KZ,N),ones(KZ,K+KZ+N*T)];
  % kappa
HPattern((6*N+2+K+KZ+1):end,:)=HPattern((6*N+2+K+KZ+1):end,:)+...
    [repmat([repmat(eye(N),1,2),ones(N,2*N+2),eye(N),zeros(N),ones(N,K+KZ)],T,1),eye(N*T)];
HPattern=1*(HPattern>0);

HPattern=sparse(HPattern);

clear n

% options
options=optimset('Display','off','JacobPattern',JPattern,'HessPattern',HPattern);


%% starting values

% "short" replication
load('robustLTnl_replication','MAP')
pars0=MAP;

% % "long" replication
% load('robustLTnl_replication','pars0')


%% maximum-a-posteriori estimator
[MAP,lp,exit]=knitromatlab(@(x)log_posterior(x,D,M,X,Z,...
    a0,om_a,th0,om_th,b0A,b0D,om_b,s_v,d_v,f0,om_f,s_vs,d_vs),pars0,[],[],...
  [],[],[-Inf(3*N+2,1);zeros(N,1);-Inf(N,1);zeros(N,1);-Inf(K,1);zeros(KZ,1);-Inf(N*T,1)],...
  Inf(length(pars0),1),@(x)eq_con_jac_NOlearn(x,y,D,M,X,Z,sgi_nodes,sgi_weights,ep_sd,Sigma),[],...
  options,'knitro.opt');


%%
clearvars -except MAP exit
save('MAProbustLTnl.mat')



